package linux;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Random;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.apache.commons.lang3.math.NumberUtils;
import org.apache.commons.net.whois.WhoisClient;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;

import jodd.io.NetUtil;

public class FetchDomainCountries {
	private static ObjectMapper objMapper = new ObjectMapper();
	private static List<Map<String, String>> tlds= new ArrayList<>();
	private static long lastQueryWhoisTime =0;
	private static int delay = 2;
	// repeat limit
	private static int repeatCounter = 5;
	private static Random rand = new Random();
	private static long timeoutMs = 5000L;
	
	static {
		try {// load all top level domains
			tlds.addAll( objMapper.readValue(ClassLoader.getSystemResourceAsStream("top_level_domains.json"), new TypeReference<List<Map<String, String>>>() {}));
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
	public static void main(String[] args) {
		// fetch linux maintainers
		String maintainerFileUrl = "https://github.com/torvalds/linux/raw/refs/heads/master/MAINTAINERS";
		List<String> maintainerFileLines = new ArrayList<>();
		try {
			maintainerFileLines.addAll(
					Arrays.asList(NetUtil.downloadString(maintainerFileUrl, StandardCharsets.UTF_8).split("\n")));
		} catch (IOException e) {
			e.printStackTrace();
		}
		
		// get domains from maintainers email  
		List<String> domains  = maintainerFileLines.stream()
				.filter(l -> l.startsWith("M:")) // get only  maintainers
				.map(l -> l.substring(2).trim())
				.distinct()
				.map(l -> { // get email
					// just "torvalds@linux-foundation.org"
					String email = l;
					if (l.contains("<") && l.contains(">")) {
						// format like "Linus Torvalds <torvalds@linux-foundation.org>"
						email = l.substring(l.indexOf("<") + 1, l.indexOf(">"));
					}
					return email;
				})
				.distinct()
				.map(l -> l.substring(l.indexOf("@") + 1))
				.distinct()
				.map(l -> {// get domain
					String[] split = l.split("\\.");
					if (split.length > 2) {
						return split[split.length - 2] + "." + split[split.length - 1];
					}
					return l;
				})
				.distinct()
				.toList();
		System.out.println("all domains:"+domains.size());
		System.out.println(String.join("\n", domains));
		Map<String, String> domainsWithCountry = getCountryFromDomains(domains);
		try {
			System.out.println("domains with country:");
			System.out.println(objMapper.writerWithDefaultPrettyPrinter().writeValueAsString(domainsWithCountry));
		} catch (JsonProcessingException e) {
			e.printStackTrace();
		}
	}
	private static Map<String, String> getCountryFromDomains(List<String> domains) {
		Map<String, String> domainsWithCountry = new HashMap<>();
		List<String> failDomains = new ArrayList<>();
		for(int i=0;i<domains.size();i++) {
			String domain = domains.get(i);
			String country = getCountryFromDomain(domain);
			if(country==null) {
				failDomains.add(domain);
			}else {
				domainsWithCountry.put(domain,   country);
			}
			System.out.printf("(%d/%d)%s\t%s%n",i+1,domains.size(),domain,country);
		}
		if(!failDomains.isEmpty()) {
			delay++;
			repeatCounter--;
			if(repeatCounter>0) {
				domainsWithCountry.putAll(getCountryFromDomains(failDomains));
				
			}else {
				domainsWithCountry.putAll(failDomains.stream().collect(Collectors.toMap(d->d, d->"UNKNOWN")));
				
			}
		}
		return domainsWithCountry;
	}
	private static String getCountryFromDomain(String domain) {
		
		// query  tld
		String tld = domain.substring(domain.lastIndexOf(".") );
		Optional<Map<String, String>> countryDomain = tlds.stream()
				.filter(t->t.get("domain").equalsIgnoreCase(tld) 
						&& t.get("type").equalsIgnoreCase("country-code"))
				.findAny();
		if(countryDomain.isPresent()){
			return countryDomain.get().get("domain").substring(1).toUpperCase();
		}
		
		String targetDomain = domain;
		String[] split = domain.split("\\.");
		if(split.length>2){
			targetDomain = split[split.length-2]+"."+split[split.length-1];
		}
		String country = null;
		// query whois server by http
		try {
			long delayMs = 1000L * Math.min(delay, 30);
			if(System.currentTimeMillis() - lastQueryWhoisTime < delayMs) {
				Thread.sleep(delayMs);
			}
			List<Function<String, String>> functions = Arrays.asList(
					d->getCountryFromDomainByWhoisCom(d)
					, d->getCountryFromDomainByTencent(d)
					, d->getCountryFromDomainByWhois365(d)
					, d->getCountryFromDomainByJsonCN(d));
			for(int i=0;i<functions.size();i++) {
				country = functions.get((i+repeatCounter)%functions.size()).apply(targetDomain);
				if(country!=null) {
					break;
				}
			}
			
		} catch (Exception e) {
			e.printStackTrace();
		}finally {
			lastQueryWhoisTime = System.currentTimeMillis();
			
		}
		if(country ==null) {
			
			// query whois server by socket
			String s = queryWhoisServer(targetDomain);
			country = getRegistrantCountryFromWhoisResult(s);
			if (country == null) {
				String whoisServer = getWhoisServerFromWhoisResult(s);
				if (whoisServer != null) {
					s = queryWhoisServer(targetDomain, whoisServer);
					country = getRegistrantCountryFromWhoisResult(s);
				}
			}
		}
				
		if (country != null) {
			return country.toUpperCase();
		}
		
		return null;
	}
	private static String getCountryFromDomainByWhoisCom(String domain) {
		String whoisUrl = "https://www.whois.com/whois/"+domain+"?"+(char)rand.nextInt(Character.MAX_VALUE)+"="+Math.random();
		try {
			HttpRequest request = HttpRequest.newBuilder(new URI(whoisUrl))
					.timeout(Duration.ofMillis(timeoutMs))
					.GET()
					.build();
			HttpResponse<String> result = HttpClient.newHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
			String body =  result.body();
			int registrarDataStart = body.toUpperCase().lastIndexOf("REGISTRANT COUNTRY");
			if(registrarDataStart>-1) {
				
				registrarDataStart = body.toUpperCase().indexOf(":", registrarDataStart) + 1;
				int registrarDataEnd = body.toUpperCase().indexOf("\n", registrarDataStart);
				return body.substring(registrarDataStart, registrarDataEnd).trim();
			}
		} catch ( Exception e) {
			e.printStackTrace();
		}
		return null;
	}
	private static String getCountryFromDomainByTencent(String domain) {
		try {

			String tencentWhoisUrl = "https://dnspod.cloud.tencent.com/cgi/capi?action=DescribeWhoisInfoSpecial&csrfCode=&innerCapiMark=1";
			HttpRequest request = HttpRequest.newBuilder(new URI(tencentWhoisUrl))
					.setHeader("Content-Type", "application/json; charset=UTF-8")
					.timeout(Duration.ofMillis(timeoutMs))
					.POST(HttpRequest.BodyPublishers.ofString("{\"Version\":\"2018-08-08\",\"serviceType\":\"domain\",\"api\":\"DescribeWhoisInfoSpecial\",\"DomainName\":\""+domain+"\",\"dpNodeCustomClientIPField\":\"RealClientIp\"}" ))
					.build();
			HttpResponse<String> result = HttpClient.newHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
			String body =  result.body();
			int registrarDataStart = body.toUpperCase().lastIndexOf("REGISTRANT COUNTRY");
			if(registrarDataStart<0) {
				return null;
			}
			registrarDataStart = body.toUpperCase().indexOf(":", registrarDataStart) + 1;
			int registrarDataEnd = body.toUpperCase().indexOf("\\", registrarDataStart);
			String country =body.substring(registrarDataStart, registrarDataEnd).trim();
		
			return country.toUpperCase();
		} catch (Exception e) {
			e.printStackTrace();
		}
		return null;
	}
	private static String getCountryFromDomainByWhois365(String domain) {
		String whoisUrl = "https://www.whois365.com/en/domain/"+domain+"?"+(char)rand.nextInt(Character.MAX_VALUE)+"="+Math.random();
		try {
			HttpRequest request = HttpRequest.newBuilder(new URI(whoisUrl))
					.timeout(Duration.ofMillis(timeoutMs))
					.GET()
					.build();
			HttpResponse<String> result = HttpClient.newHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
			String body =  result.body();
			int registrarDataStart = body.toUpperCase().lastIndexOf("REGISTRANT COUNTRY");
			if(registrarDataStart>-1) {
				
				registrarDataStart = body.toUpperCase().indexOf(":", registrarDataStart) + 1;
				int registrarDataEnd = body.toUpperCase().indexOf("<BR", registrarDataStart);
				return body.substring(registrarDataStart, registrarDataEnd).trim();
			}
		} catch ( Exception e) {
			e.printStackTrace();
		}
		return null;
	}
	private static String getCountryFromDomainByJsonCN(String domain) {
		String whoisUrl = "http://www.jsons.cn/whois/?txt_url="+domain+"&"+(char)rand.nextInt(Character.MAX_VALUE)+"="+Math.random();
		try {
			HttpRequest request = HttpRequest.newBuilder(new URI(whoisUrl))
					.timeout(Duration.ofMillis(timeoutMs))
					.GET()
					.build();
			HttpResponse<String> result = HttpClient.newHttpClient().send(request, HttpResponse.BodyHandlers.ofString());
			String body =  result.body();
			int registrarDataStart = body.toUpperCase().lastIndexOf("REGISTRANT&NBSP;COUNTRY");
			if(registrarDataStart>-1) {
				
				registrarDataStart = body.toUpperCase().indexOf(":", registrarDataStart) + 1;
				int registrarDataEnd = body.toUpperCase().indexOf("\n", registrarDataStart);
				return body.substring(registrarDataStart, registrarDataEnd).replace("&nbsp;", "").trim();
			}
		} catch ( Exception e) {
			e.printStackTrace();
		}
		return null;
	}

	private static String queryWhoisServer(String domain, String... whoisServer) {
		List<String> whoisServers = new ArrayList<>();
		if (whoisServer.length == 0) {
			whoisServers.add(WhoisClient.DEFAULT_HOST);
		} else {
			whoisServers.addAll(Arrays.asList(whoisServer));
		}
		for (String server : whoisServers) {
			try {
				WhoisClient whois = new WhoisClient();
				whois.setConnectTimeout((int)timeoutMs);
				// get proxy variables in environment
				String proxyHost = System.getenv("SOCKS_PROXY_HOST");
				String socksProxyPortEnv = System.getenv("SOCKS_PROXY_PORT");
				int proxyPort =  NumberUtils.isParsable(socksProxyPortEnv)?Double.valueOf(socksProxyPortEnv).intValue():-1;
				if (proxyHost != null && proxyPort > 0) {
					whois.setProxy(new Proxy(Proxy.Type.SOCKS, new InetSocketAddress(proxyHost, proxyPort)));
				}
				
				whois.connect(server);
				String s = whois.query(domain);
				whois.disconnect();
				return s;
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
		return null;
	}

	private static String getWhoisServerFromWhoisResult(String s) {
		if (s == null) {
			return null;
		}
		Optional<String> whoisServerLine = Arrays.stream(s.split("\n"))
				.filter(l -> l.trim().toUpperCase().startsWith("REGISTRAR WHOIS SERVER:")).findFirst();
		if (whoisServerLine.isPresent()) {
			return whoisServerLine.get().substring(whoisServerLine.get().indexOf(":") + 1).trim();
		}
		return null;
	}

	private static String getRegistrantCountryFromWhoisResult(String s) {
		if (s == null) {
			return null;
		}
		Optional<String> registrantCountryLine = Arrays.stream(s.split("\n"))
				.filter(l -> l.trim().toUpperCase().startsWith("REGISTRANT COUNTRY:")).findFirst();
		if (registrantCountryLine.isPresent()) {
			return registrantCountryLine.get().substring(registrantCountryLine.get().indexOf(":") + 1).trim();
		}
		return null;
	}
}
